import os

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image

from model import CNN

# ===== 初始化模型 =====
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CNN().to(device)
model.load_state_dict(torch.load("best_model.pth", map_location=device))
model.eval()

def segment_characters(img_array):
    # 二值化处理
    _, thresh = cv2.threshold(img_array, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # 反转图像，白字黑底
    thresh = 255 - thresh

    # 膨胀处理连接数字碎片
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    dilated = cv2.dilate(thresh, kernel, iterations=1)

    # 查找轮廓
    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    char_regions = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        if w > 5 and h > 10:  # 简单过滤噪声
            char_regions.append((x, y, w, h))

    # 左到右排序
    char_regions.sort(key=lambda b: b[0])

    char_images = []
    for (x, y, w, h) in char_regions:
        char_img = img_array[y:y + h, x:x + w]
        char_images.append(char_img)

    return char_images

# ===== 图像预处理相关函数 =====
def dynamic_color_normalization(img_array):
    if len(img_array.shape) == 3:
        lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
        img_array = lab[:, :, 0]
    _, thresh = cv2.threshold(img_array, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    white_pixels = np.sum(thresh == 255)
    black_pixels = np.sum(thresh == 0)
    if white_pixels > black_pixels:
        img_array = 255 - img_array
    return img_array


def enhance_digit_region(img_array):
    edges = cv2.Canny(img_array, 50, 150)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    closed = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
    contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        max_contour = max(contours, key=cv2.contourArea)
        x, y, w, h = cv2.boundingRect(max_contour)
        margin = int(max(w, h) * 0.2)
        x = max(0, x - margin)
        y = max(0, y - margin)
        w = min(img_array.shape[1] - x, w + 2 * margin)
        h = min(img_array.shape[0] - y, h + 2 * margin)
        roi = img_array[y:y + h, x:x + w]
        scaled = cv2.resize(roi, (20, 20))
        img_array = np.zeros((28, 28), dtype=np.uint8)
        img_array[4:24, 4:24] = scaled
    return img_array


def contrast_enhancement(img_array):
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(4, 4))
    img_array = clahe.apply(img_array)
    p2, p98 = np.percentile(img_array, (2, 98))
    img_array = np.clip((img_array - p2) * 255.0 / (p98 - p2), 0, 255).astype(np.uint8)
    gamma = 1.5 if np.mean(img_array) < 127 else 0.7
    img_array = np.power(img_array / 255.0, gamma) * 255.0
    return img_array.astype(np.uint8)


def preprocess_image(img_array):
    img_array = dynamic_color_normalization(img_array)
    img_array = enhance_digit_region(img_array)
    img_array = cv2.resize(img_array, (28, 28))
    img_array = contrast_enhancement(img_array)
    img_tensor = torch.from_numpy(img_array).float() / 255.0
    img_tensor = img_tensor.unsqueeze(0).unsqueeze(0)
    return img_tensor.to(device)


# ===== 主预测函数 =====
def predict_images_in_folder(folder_path, model_path="best_model.pth"):
    if not os.path.exists(model_path):
        print(f"❌ 模型文件未找到: {model_path}")
        return

    if not os.path.exists(folder_path):
        print(f"❌ 文件夹不存在：{folder_path}")
        return

    image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    if not image_files:
        print("⚠️ 没有找到任何图片")
        return

    for filename in sorted(image_files):
        img_path = os.path.join(folder_path, filename)
        try:
            # 加载图像（灰度）
            img = Image.open(img_path).convert('L')
            img_array = np.array(img)

            # 分割字符
            char_images = segment_characters(img_array)
            if not char_images:
                print(f"⚠️ 没有检测到字符：{filename}")
                continue

            print(f"🖼️ {filename} ➜ 检测到 {len(char_images)} 个字符")

            for i, char_img in enumerate(char_images):
                # 预处理单个字符
                char_tensor = preprocess_image(Image.fromarray(char_img))

                # 模型预测
                with torch.no_grad():
                    output = model(char_tensor)
                    _, predicted = torch.max(output.data, 1)
                    pred_label = predicted.item()

                print(f"   🔢 第 {i+1} 个字符预测: {pred_label}")
                plt.imshow(char_tensor.squeeze().cpu(), cmap='gray')
                plt.title(f"{filename} - Char {i+1}: {pred_label}")
                plt.axis('off')
                plt.show()
        except Exception as e:
            print(f"跳过图像 {filename}（处理失败: {str(e)}）")


# ===== 启动测试 =====
if __name__ == "__main__":
    predict_images_in_folder("test_images")
